import cv2
import numpy as np
from collections import Counter
import pickle


# ------ 将图片进行灰度化处理 --------
def rgb2gray(img):
    imgGray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)

    # ------ 查看灰度图片 -------
    # cv2.imshow('imgGray', imgGray)

    # ------ 查看灰度图片形状 --------
    # print("imgShape:", imgGray.shape)
    return imgGray


# ------ 将图片进行灰度化处理 --------
def rgb2hsv(img):
    # ------ 对HSV处理后的图片 --------
    imgHSV = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # ------ 对HSV图片 -------
    cv2.imshow('imgHSV', imgHSV)

    return imgHSV


# ------ 多维数组拉平一维 --------
def flatten2one(img):
    oneDimension = img.flatten()
    return oneDimension


# https://www.cnblogs.com/ssyfj/p/9272615.html 二值化处理参考

# ------ 将图片进行二值化处理 --------
def binary_average(imgGray, oneDimension):
    r, b = cv2.threshold(imgGray, sum(oneDimension) / oneDimension.size, 255, cv2.THRESH_BINARY)

    return r, b


# ------ 将图片进行二值化处理(大津法) --------
def binary_ostu(imgGray):
    r, b = cv2.threshold(imgGray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
    # cv2.imshow('binary_ostu_img', b)
    return r, b

def get_plantZone(gmm,initImg,L,W):
    newImg = initImg.reshape(L * W, 3)

    gmmModel = gmm.predict(newImg)
    for i in range(gmmModel.size):
        if (gmmModel[i] == 0):
            gmmModel[i] = 255

    gmmModel = gmmModel.reshape((L, W))

    return np.array(gmmModel, np.uint8)

# ------ 打印干湿比 --------
def print_ratio_percent(b, i=666):
    params = Counter(b.flatten())
    print("去除植被前：data_{}.jpg 干湿比为：{}, 湿润点比例为：{}%".format(i, params[255] / params[0], params[0] / (params[0] + params[255])))


# ------ 打印经过植被去除后的干湿比、含水量 --------
def print_ratio_percent_remove_plant(b, gmmModel, index=666):
    # 干燥点
    x = 0
    # 湿润点
    y = 0
    for i in range(gmmModel.shape[0]):
        for j in range(gmmModel.shape[0]):
            if gmmModel[i][j] == 1:
                if b[i][j] == 255:
                    x = x + 1
                else:
                    y = y + 1

    print("去除植被后：data_{}.jpg 干湿比为：{}, 湿润点比例为：{}%".format(index, x / y, y / (x + y)))


if __name__ == "__main__":
    # ------ 设置打印是完整打印 --------
    np.set_printoptions(threshold=np.inf)

    with open('./model/model_0.bin', 'rb') as f:
        gmm = pickle.loads(f.read())

    # 定义读取图像
    initImg = cv2.imread("./data/data_24.jpg")
    # initImg = cv2.imread("./data/ningxiang/2.jpg")

    cv2.imshow('initImg', initImg)

    L,W = initImg.shape[0], initImg.shape[1]

    gmmModel = get_plantZone(gmm,initImg,L,W)
    cv2.imshow('plant', gmmModel)

    # 灰度化图像
    imgGray = rgb2gray(initImg)

    # r: 阈值, b: 二值化图像, 大津法
    r, b = binary_ostu(imgGray)

    cv2.imshow('Binary', b)

    # 去除植被前
    print_ratio_percent(b)
    # 去除植被后
    print_ratio_percent_remove_plant(b, gmmModel)

    cv2.waitKey(0)
    cv2.destroyAllWindows()
